#!/usr/bin/env python
from utils import GServer, GClient, config, GAnswer
from mime import get_param, convert_to_b64, structure, convert_to_log, message_from_string, save_message
from smtplib import CRLF, quotedata, quoteaddr
import base64, time

class SmtpClient(GClient):
	def send_cmd(self, cmd):
		a = self.send(cmd + '\r\n')
		if not a.ok:
			return a
		a = self.recv()
		if not a.ok:
			return a
		j = a.str
		if j[:3].isdigit():
			rc = int(j[:3])
		else:
			rc = 0
		return GAnswer(True, repr(j), rc = rc)

	def connect(self):
		a = GClient.connect(self)
		if not a.ok: return a
		return self.recv()

	def helo(self):
		return self.send_cmd('HELO 127.0.0.1')

	def ehlo(self):
		return self.send_cmd('EHLO 127.0.0.1')

	def quit(self):
		return self.send_cmd('QUIT')

	def login(self, user, password):
		return self.send_cmd('AUTH PLAIN ' + base64.b64encode('\0%s\0%s' % (user, password)))

	def send_msg_data(self, obj):
		return self.send_msg(obj)

	def send_msg_bdat(self, obj, chunk_size = 8192):
		return self.send_msg(obj, True, chunk_size)

	def send_msg(self, obj, use_bdat = False, dbat_size = 8192):
		msg = config.get_obj(obj)

		if msg is None:
			return GAnswer(False, 'No such message: %s' % obj)
		from_addr = msg.get_param('from', False)
		to_addr = '%s;%s;%s' % (msg.get_param('to', False), msg.get_param('cc', False), msg.get_param('bcc', False))
		a = self.send_cmd('MAIL FROM: %s' % quoteaddr(from_addr))
		self.log('MAIL FROM: %s' % quoteaddr(from_addr))
		self.loga(a)
		l = to_addr.split(';')
		for i in l:
			if i == 'NIL' or len(i) == 0: continue
			if i[0] == ' ': i = i[1:]
			a = self.send_cmd('RCPT TO: %s' % quoteaddr(i))
			self.log('RCPT TO: %s' % quoteaddr(i))
			self.loga(a)

		if use_bdat:			
			p = msg.as_string(skip_bcc = True)
			self.log('sending BDAT msg:%s len:%i' % (obj, len(p)))
			chunks = len(p) / dbat_size
			from cStringIO import StringIO
			encoded_fp = StringIO(p)
			chunk = encoded_fp.read(dbat_size)
			a = None
			i = 0
			while chunk:
				if len(chunk) < dbat_size:
					self.send('BDAT %i LAST%s' % (len(chunk), CRLF))
				else:
					self.send('BDAT %i%s' % (len(chunk), CRLF))
				self.send(chunk)
				chunk = encoded_fp.read(dbat_size)
				a = self.recv()
			return a
		else:
			a = self.send_cmd('DATA')
			q = quotedata(msg.as_string(skip_bcc = True))
			if q[-2:] != CRLF: q += CRLF
			a = self.send_text(q)
			self.loga(a)
			a = self.send_cmd('.')
			return a

	def starttls(self):
		a = self.send_cmd('STARTTLS')
		if a.rc == 220:
			self.ssl()
		else:
			a.ok = False
		return a

#
# This drives the state for a single RFC821 message.
# Author: L. Smithson (lsmithson@open-networks.co.uk)
# + some new things
#
class SMTPServerEngine:
	"""
	Server engine that calls methods on the SMTPServerInterface object
	passed at construction time. It is constructed with a bound socket
	connection to a client. The 'chug' method drives the state,
	returning when the client RFC821 transaction is complete. 
	"""

	ST_INIT = 0
	ST_HELO = 1
	ST_MAIL = 2
	ST_RCPT = 3
	ST_DATA = 4
	ST_BDAT = 5
	ST_BDAT_LAST = 6
	ST_QUIT = 7
	
	greeting = "220 Welcome\r\n"
	def __init__(self, socket, impl):
		self.impl = impl;
		self.socket = socket;
		self.state = SMTPServerEngine.ST_INIT
		
	def chug(self):
		"""
		Chug the engine, till QUIT is received from the client. As
		each RFC821 message is received, calls are made on the
		SMTPServerInterface methods on the object passed at
		construction time.
		"""
		
		self.socket.send(self.greeting)
		bdata = ''
		while True:
			data = ''
			if self.state == SMTPServerEngine.ST_BDAT or self.state == SMTPServerEngine.ST_BDAT_LAST:
				bdata += self.socket.recv(self.bdat_size)
				if self.state == SMTPServerEngine.ST_BDAT_LAST:
					self.state = SMTPServerEngine.ST_HELO
					rv = self.impl.bdat(bdata)
					self.socket.send(rv + CRLF)
			completeLine = 0
			# Make sure an entire line is received before handing off
			# to the state engine. Thanks to John Hall for pointing
			# this out.
			while not completeLine:
				lump = self.socket.recv(1024);
				if len(lump):
					data += lump
					if self.state == SMTPServerEngine.ST_BDAT:
						nl = data.find(CRLF) + 2
						bdata += data[nl:]
						data = data[:nl]
					if (len(data) >= 2) and data[-2:] == '\r\n':
						completeLine = 1
						if self.state != SMTPServerEngine.ST_DATA:
							rsp, keep = self.doCommand(data)
						else:
							rsp = self.doData(data)
							if rsp is None:
								continue
						self.socket.send(rsp + "\r\n")
						if keep == 0:
							self.socket.close()
							return
				else:
					# EOF
					return
		return
			
	def doCommand(self, data, rest = None):
		"""Process a single SMTP Command"""
		cmd = data[0:4].upper()
		keep = 1
		rv = None
		if cmd == "HELO":
			self.state = SMTPServerEngine.ST_HELO
			rv = self.impl.helo(data)
		elif cmd == "EHLO":
			self.state = SMTPServerEngine.ST_HELO
			rv = self.impl.ehlo(data)
		elif cmd == "AUTH":
			self.state = SMTPServerEngine.ST_HELO
			rv = self.impl.auth(data)
		elif cmd == "RSET":
			rv = self.impl.reset(data)
			self.dataAccum = ""
			self.state = SMTPServerEngine.ST_INIT
		elif cmd == "NOOP":
			pass
		elif cmd == "QUIT":
			rv = self.impl.quit(data)
			keep = 0
		elif cmd == "MAIL":
			if self.state != SMTPServerEngine.ST_HELO:
				return ("503 Bad command sequence", 1)
			self.state = SMTPServerEngine.ST_MAIL
			rv = self.impl.mailFrom(data)
		elif cmd == "RCPT":
			if (self.state != SMTPServerEngine.ST_MAIL) and (self.state != SMTPServerEngine.ST_RCPT):
				return ("503 Bad command sequence", 1)
			self.state = SMTPServerEngine.ST_RCPT
			rv = self.impl.rcptTo(data)
		elif cmd == "DATA":
			if self.state != SMTPServerEngine.ST_RCPT:
				return ("503 Bad command sequence", 1)
			self.state = SMTPServerEngine.ST_DATA
			self.dataAccum = ""
			return ("354 OK, Enter data, terminated with a \\r\\n.\\r\\n", 1)
		elif cmd == "BDAT":
			if self.state != SMTPServerEngine.ST_RCPT and self.state != SMTPServerEngine.ST_BDAT:
				return ("503 Bad command sequence", 1)
			s = data.split()
			self.bdat_size = int(s[1])
			if len(s) > 2 and s[2].upper() == 'LAST':
				self.state = SMTPServerEngine.ST_BDAT_LAST
			else:
				self.state = SMTPServerEngine.ST_BDAT
		else:
			return ("505 Eh? WTF was that?", 1)

		if rv:
			return (rv, keep)
		else:
			return("250 OK", keep)

	def doData(self, data):
		"""
		Process SMTP Data. Accumulates client DATA until the
		terminator is found.
		"""
		self.dataAccum = self.dataAccum + data
		if len(self.dataAccum) > 4 and self.dataAccum[-5:] == '\r\n.\r\n':
			self.dataAccum = self.dataAccum[:-5]
			rv = self.impl.data(self.dataAccum)
			self.state = SMTPServerEngine.ST_HELO
			if rv:
				return rv
			else:
				return "250 OK - Data and terminator. found"
		else:
			return None

class SmtpServer(GServer):
	def start(self):
		return GServer.start(self)

	def handler(self, socket, address):
		engine = SMTPServerEngine(socket, self)
		engine.greeting = self.welcome_resp
		engine.chug()

	def helo(self, args):
		self.log('Received "helo" %s' % args[:-1])
		return self.helo_resp

	def ehlo(self, args):
		self.log('Received "ehlo" %s' % args[:-1])
		return self.ehlo_resp

	def auth(self, args):
		a = args.split()
		self.log('Received "AUTH" %s' % args[:-1])
		self.log('user/pw: %s' % base64.b64decode(a[2])[1:].split('\00'))
		#return "334 ok. go on."
		return self.auth_resp
    
	def mailFrom(self, args):
		self.log('Received "MAIL FROM:" %s' % args[:-1])

	def rcptTo(self, args):
		self.log('Received "RCPT TO:" %s' % args[:-1])

	def data(self, args):
		self.log('Received "DATA"')
		return self._data(args)

	def bdat(self, args):
		self.log('Received "BDAT"')
		return self._data(args)

	def _data(self, args):
		m = message_from_string(args)
		if self.save:
			save_message(m, '%s/%s/%s' % (config.main.work_dir, self.name, str(time.time())))
		convert_to_log(m)
		self.log('message:\n' + m.as_string())
		self.log('struct:\n' + structure(m))
		return '250 OK'

	def quit(self, args):
		self.log('Received "QUIT" %s' % args[:-1])
		return '221 Goodbye'

	def reset(self, args):
		self.log('Received "RSET" %s' % args[:-1])

